// Copyright 2017 The Ray Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "ray/core_worker/context.h"

#include <google/protobuf/util/json_util.h>

#include <memory>
#include <string>

#include "ray/common/runtime_env_common.h"

namespace ray {
namespace core {
namespace {
const rpc::JobConfig kDefaultJobConfig{};
}  // namespace

/// per-thread context for core worker.
struct WorkerThreadContext {
  explicit WorkerThreadContext(const JobID &job_id)
      : max_num_generator_returns_(RayConfig::instance().max_num_generator_returns()) {
    SetCurrentTaskId(TaskID::FromRandom(job_id), /*attempt_number=*/0);
  }

  uint64_t GetNextTaskIndex() { return ++task_index_; }

  uint64_t GetTaskIndex() const { return task_index_; }

  /// Returns the next put object index. The index starts at the number of
  /// return values for the current task in order to keep the put indices from
  /// conflicting with return object indices. 1 <= idx <= NumReturns() is reserved for
  /// return objects, while idx > NumReturns is available for put objects.
  ObjectIDIndexType GetNextPutIndex() {
    // If current_task_ is nullptr, we assume that we're in the event loop thread and
    // are executing async tasks; in this case, we're using a fake, random task ID
    // for put objects, so there's no risk of creating put object IDs that conflict with
    // return object IDs (none of the latter are created). The put counter will never
    // reset and will therefore continue to increment for the lifetime of the event
    // loop thread (ResetCurrentTask and SetCurrenTask will never be called in the
    // thread), so there's no risk of conflicting put object IDs, either.
    // See https://github.com/ray-project/ray/issues/10324 for further details.
    auto num_returns = current_task_ != nullptr ? current_task_->NumReturns() : 0;

    // We reserve max_num_generator_returns_ number of indexes for the generator
    // return so that all generator return can have consistent ids given an index.
    return num_returns + max_num_generator_returns_ + ++put_counter_;
  }

  const TaskID &GetCurrentTaskID() const { return current_task_id_; }

  std::shared_ptr<const TaskSpecification> GetCurrentTask() const {
    return current_task_;
  }

  void SetCurrentTaskId(const TaskID &task_id, uint64_t attempt_number) {
    current_task_id_ = task_id;
    if (!current_task_id_.IsNil()) {
      current_internal_task_id_ = TaskID::ForExecutionAttempt(task_id, attempt_number);
    } else {
      current_internal_task_id_ = TaskID::Nil();
    }
  }

  const TaskID &GetCurrentInternalTaskId() const { return current_internal_task_id_; }

  const PlacementGroupID &GetCurrentPlacementGroupId() const {
    return current_placement_group_id_;
  }

  void SetCurrentPlacementGroupId(const PlacementGroupID &placement_group_id) {
    current_placement_group_id_ = placement_group_id;
  }

  void SetPlacementGroupCaptureChildTasks(bool placement_group_capture_child_tasks) {
    placement_group_capture_child_tasks_ = placement_group_capture_child_tasks;
  }

  bool PlacementGroupCaptureChildTasks() const {
    return placement_group_capture_child_tasks_;
  }

  void SetCurrentTask(const TaskSpecification &task_spec) {
    RAY_CHECK(task_index_ == 0);
    RAY_CHECK(put_counter_ == 0);
    SetCurrentTaskId(task_spec.TaskId(), task_spec.AttemptNumber());
    SetCurrentPlacementGroupId(task_spec.PlacementGroupBundleId().first);
    SetPlacementGroupCaptureChildTasks(task_spec.PlacementGroupCaptureChildTasks());
    current_task_ = std::make_shared<const TaskSpecification>(task_spec);
  }

  void ResetCurrentTask() {
    SetCurrentTaskId(TaskID::Nil(), /*attempt_number=*/0);
    task_index_ = 0;
    put_counter_ = 0;
  }

  uint32_t GetMaxNumGeneratorReturnIndex() const { return max_num_generator_returns_; }

 private:
  /// The task ID for current task.
  TaskID current_task_id_;

  /// The internal task ID, used to seed any TaskIDs or ObjectIDs generated by
  /// the current task (via ray.put or .remote calls). As tasks can be executed
  /// multiple times, this ID represents the n-th execution of the current
  /// task.
  /// See https://github.com/ray-project/ray/issues/20713.
  TaskID current_internal_task_id_;

  /// The current task.
  std::shared_ptr<const TaskSpecification> current_task_;

  /// Number of tasks that have been submitted from current task.
  uint64_t task_index_{0};

  static_assert(sizeof(task_index_) == TaskID::Size() - ActorID::Size(),
                "Size of task_index_ doesn't match the unique bytes of a TaskID.");

  /// A running counter for the number of object puts carried out in the current task.
  /// Used to calculate the object index for put object ObjectIDs.
  ObjectIDIndexType put_counter_{0};

  static_assert(sizeof(put_counter_) == ObjectID::Size() - TaskID::Size(),
                "Size of put_counter_ doesn't match the unique bytes of an ObjectID.");

  /// Placement group id that the current task belongs to.
  /// NOTE: The top level `WorkerContext` will also have placement_group_id
  ///   which is set when actors are created. It is because we'd like to keep track
  ///   thread local placement group id for tasks, and the process placement group id for
  ///   actors.
  PlacementGroupID current_placement_group_id_;

  /// Whether or not child tasks are captured in the parent's placement group implicitly.
  bool placement_group_capture_child_tasks_ = false;

  /// The maximum number of generator return values.
  uint32_t max_num_generator_returns_;
};

thread_local std::unique_ptr<WorkerThreadContext> WorkerContext::thread_context_ =
    nullptr;

WorkerContext::WorkerContext(WorkerType worker_type,
                             const WorkerID &worker_id,
                             const JobID &job_id)
    : worker_type_(worker_type),
      worker_id_(worker_id),
      current_job_id_(job_id),
      current_actor_id_(ActorID::Nil()),
      current_actor_placement_group_id_(PlacementGroupID::Nil()),
      placement_group_capture_child_tasks_(false),
      main_thread_id_(boost::this_thread::get_id()),
      root_detached_actor_id_(ActorID::Nil()) {
  // For worker main thread which initializes the WorkerContext,
  // set task_id according to whether current worker is a driver.
  // (For other threads it's set to random ID via GetThreadContext).
  if (worker_type_ == WorkerType::DRIVER) {
    RAY_CHECK(!current_job_id_.IsNil());
    GetThreadContext().SetCurrentTaskId(TaskID::ForDriverTask(job_id),
                                        /*attempt_number=*/0);
    // Driver runs in the main thread.
    {
      absl::WriterMutexLock lock(&mutex_);
      main_thread_or_actor_creation_task_id_ = TaskID::ForDriverTask(job_id);
    }
  }
}

WorkerType WorkerContext::GetWorkerType() const { return worker_type_; }

const WorkerID &WorkerContext::GetWorkerID() const { return worker_id_; }

uint64_t WorkerContext::GetNextTaskIndex() {
  return GetThreadContext().GetNextTaskIndex();
}

uint64_t WorkerContext::GetTaskIndex() const { return GetThreadContext().GetTaskIndex(); }

ObjectIDIndexType WorkerContext::GetNextPutIndex() {
  return GetThreadContext().GetNextPutIndex();
}

void WorkerContext::MaybeInitializeJobInfo(const JobID &job_id,
                                           const rpc::JobConfig &job_config) {
  {
    absl::ReaderMutexLock lock(&mutex_);
    if (!current_job_id_.IsNil() && job_config_.has_value()) {
      RAY_CHECK(current_job_id_ == job_id);
      return;
    }
  }
  absl::WriterMutexLock lock(&mutex_);
  current_job_id_ = job_id;
  job_config_ = job_config;
  RAY_CHECK(current_job_id_ == job_id);
}

int64_t WorkerContext::GetTaskDepth() const {
  absl::ReaderMutexLock lock(&mutex_);
  return task_depth_;
}

JobID WorkerContext::GetCurrentJobID() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_job_id_;
}

rpc::JobConfig WorkerContext::GetCurrentJobConfig() const {
  absl::ReaderMutexLock lock(&mutex_);
  return job_config_.has_value() ? job_config_.value() : kDefaultJobConfig;
}

const TaskID &WorkerContext::GetCurrentTaskID() const {
  return GetThreadContext().GetCurrentTaskID();
}

const TaskID &WorkerContext::GetCurrentInternalTaskId() const {
  return GetThreadContext().GetCurrentInternalTaskId();
}

PlacementGroupID WorkerContext::GetCurrentPlacementGroupId() const {
  absl::ReaderMutexLock lock(&mutex_);
  // If the worker is an actor, we should return the actor's placement group id.
  if (current_actor_id_ != ActorID::Nil()) {
    return current_actor_placement_group_id_;
  } else {
    return GetThreadContext().GetCurrentPlacementGroupId();
  }
}

bool WorkerContext::ShouldCaptureChildTasksInPlacementGroup() const {
  absl::ReaderMutexLock lock(&mutex_);
  // If the worker is an actor, we should return the actor's placement group id.
  if (current_actor_id_ != ActorID::Nil()) {
    return placement_group_capture_child_tasks_;
  } else {
    return GetThreadContext().PlacementGroupCaptureChildTasks();
  }
}

std::shared_ptr<rpc::RuntimeEnvInfo> WorkerContext::GetCurrentRuntimeEnvInfo() const {
  absl::ReaderMutexLock lock(&mutex_);
  return runtime_env_info_;
}

const std::string &WorkerContext::GetCurrentSerializedRuntimeEnv() const {
  absl::ReaderMutexLock lock(&mutex_);
  return runtime_env_info_->serialized_runtime_env();
}

std::shared_ptr<nlohmann::json> WorkerContext::GetCurrentRuntimeEnv() const {
  absl::ReaderMutexLock lock(&mutex_);
  return runtime_env_;
}

void WorkerContext::SetCurrentTaskId(const TaskID &task_id, uint64_t attempt_number) {
  GetThreadContext().SetCurrentTaskId(task_id, attempt_number);
}

void WorkerContext::SetCurrentActorId(const ActorID &actor_id)
    ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::WriterMutexLock lock(&mutex_);
  if (!current_actor_id_.IsNil()) {
    RAY_CHECK(current_actor_id_ == actor_id);
    return;
  }
  current_actor_id_ = actor_id;
}

void WorkerContext::SetTaskDepth(int64_t depth) { task_depth_ = depth; }

void WorkerContext::SetCurrentTask(const TaskSpecification &task_spec) {
  GetThreadContext().SetCurrentTask(task_spec);

  const auto &serialized_runtime_env =
      task_spec.GetMessage().runtime_env_info().serialized_runtime_env();

  absl::WriterMutexLock lock(&mutex_);
  SetTaskDepth(task_spec.GetDepth());
  if (CurrentThreadIsMain()) {
    main_thread_or_actor_creation_task_id_ = task_spec.TaskId();
  }
  RAY_CHECK(current_job_id_ == task_spec.JobId());
  if (task_spec.IsNormalTask()) {
    current_task_is_direct_call_ = true;
    root_detached_actor_id_ = task_spec.RootDetachedActorId();
  } else if (task_spec.IsActorCreationTask()) {
    if (!current_actor_id_.IsNil()) {
      RAY_CHECK(current_actor_id_ == task_spec.ActorCreationId());
    }
    current_actor_id_ = task_spec.ActorCreationId();
    current_actor_should_exit_ = false;
    current_actor_is_direct_call_ = true;
    current_actor_max_concurrency_ = task_spec.MaxActorConcurrency();
    current_actor_is_asyncio_ = task_spec.IsAsyncioActor();
    is_detached_actor_ = task_spec.IsDetachedActor();
    current_actor_placement_group_id_ = task_spec.PlacementGroupBundleId().first;
    placement_group_capture_child_tasks_ = task_spec.PlacementGroupCaptureChildTasks();
    root_detached_actor_id_ = task_spec.RootDetachedActorId();
  } else if (task_spec.IsActorTask()) {
    RAY_CHECK(current_actor_id_ == task_spec.ActorId());
  } else {
    RAY_CHECK(false);
  }

  if (task_spec.IsNormalTask() || task_spec.IsActorCreationTask()) {
    const bool is_first_time_assignment = runtime_env_info_ == nullptr;

    // Only perform heavy-loaded assignment and parsing on first access.
    // All threads are requesting for the same parsed json result, so ok to place in
    // critical section.
    if (is_first_time_assignment) {
      runtime_env_info_ = std::make_shared<rpc::RuntimeEnvInfo>();
      *runtime_env_info_ = task_spec.RuntimeEnvInfo();

      RAY_CHECK(serialized_runtime_env_.empty());
      RAY_CHECK(runtime_env_ == nullptr);
      if (!IsRuntimeEnvEmpty(serialized_runtime_env)) {
        runtime_env_ = std::make_shared<nlohmann::json>();
        *runtime_env_ = nlohmann::json::parse(serialized_runtime_env);
      }
      serialized_runtime_env_ = serialized_runtime_env;
      return;
    }

    // Ray currently doesn't reuse worker to run tasks or actors with different runtime
    // envs.
    RAY_CHECK_EQ(serialized_runtime_env_, serialized_runtime_env);
    if (!IsRuntimeEnvEmpty(serialized_runtime_env)) {
      RAY_CHECK(runtime_env_ != nullptr);
    }
  }
}

void WorkerContext::ResetCurrentTask() { GetThreadContext().ResetCurrentTask(); }

/// NOTE: This method can't be used in fiber/async actor context.
std::shared_ptr<const TaskSpecification> WorkerContext::GetCurrentTask() const {
  return GetThreadContext().GetCurrentTask();
}

// TODO(dayshah): Fixing thread-safety-reference-return here causes Java test failures.
// Fix in follow up.
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wunknown-warning-option"
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wthread-safety-reference-return"
#endif
const ActorID &WorkerContext::GetCurrentActorID() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_actor_id_;
}
#ifdef __clang__
#pragma clang diagnostic pop
#pragma clang diagnostic pop
#endif

ActorID WorkerContext::GetRootDetachedActorID() const {
  absl::ReaderMutexLock lock(&mutex_);
  return root_detached_actor_id_;
}

bool WorkerContext::CurrentThreadIsMain() const {
  return boost::this_thread::get_id() == main_thread_id_;
}

TaskID WorkerContext::GetMainThreadOrActorCreationTaskID() const {
  absl::ReaderMutexLock lock(&mutex_);
  return main_thread_or_actor_creation_task_id_;
}

bool WorkerContext::ShouldReleaseResourcesOnBlockingCalls() const {
  // Check if we need to release resources when we block:
  //  - Driver doesn't acquire resources and thus doesn't need to release.
  //  - We only support lifetime resources for actors, which can be
  //    acquired when the actor is created, per call resources are not supported,
  //    thus we don't need to release resources for actor calls.
  return worker_type_ != WorkerType::DRIVER && !CurrentActorIsDirectCall() &&
         CurrentThreadIsMain();
}

// TODO(edoakes): simplify these checks now that we only support direct call mode.
bool WorkerContext::CurrentActorIsDirectCall() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_actor_is_direct_call_;
}

bool WorkerContext::CurrentTaskIsDirectCall() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_task_is_direct_call_ || current_actor_is_direct_call_;
}

int WorkerContext::CurrentActorMaxConcurrency() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_actor_max_concurrency_;
}

bool WorkerContext::CurrentActorIsAsync() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_actor_is_asyncio_;
}

void WorkerContext::SetCurrentActorShouldExit() {
  absl::WriterMutexLock lock(&mutex_);
  RAY_CHECK(!current_actor_id_.IsNil())
      << "SetCurrentActorShouldExit should only be used inside actors";
  current_actor_should_exit_ = true;
}

bool WorkerContext::GetCurrentActorShouldExit() const {
  absl::ReaderMutexLock lock(&mutex_);
  return current_actor_should_exit_;
}

bool WorkerContext::CurrentActorDetached() const {
  absl::ReaderMutexLock lock(&mutex_);
  return is_detached_actor_;
}

ObjectID WorkerContext::GetGeneratorReturnId(const TaskID &task_id,
                                             std::optional<ObjectIDIndexType> put_index) {
  TaskID current_task_id;
  // We only allow to specify both task id and put index or not specifying both.
  RAY_CHECK((task_id.IsNil() && !put_index.has_value()) ||
            (!task_id.IsNil() || put_index.has_value()));
  if (task_id.IsNil()) {
    const auto &task_spec = GetCurrentTask();
    current_task_id = task_spec->TaskId();
  } else {
    current_task_id = task_id;
  }

  ObjectIDIndexType current_put_index = 0;
  if (!put_index.has_value()) {
    current_put_index = GetNextPutIndex();
  } else {
    // Streaming generator case.
    current_put_index = put_index.value();
    // We don't allow to return more than GetMaxNumGeneratorReturnIndex()
    // return values.
    auto max_generator_returns = GetThreadContext().GetMaxNumGeneratorReturnIndex();
    if (put_index > max_generator_returns) {
      RAY_LOG(FATAL).WithField(current_task_id)
          << "The generator returns " << current_put_index
          << " items, which exceed the maximum number of return values allowed, "
          << max_generator_returns;
    }
  }

  return ObjectID::FromIndex(current_task_id, current_put_index);
}

/// NOTE: This method can't be used in fiber/async actor context.
WorkerThreadContext &WorkerContext::GetThreadContext() const {
  if (thread_context_ == nullptr) {
    absl::ReaderMutexLock lock(&mutex_);
    RAY_CHECK(!current_job_id_.IsNil())
        << "can't access thread context when job_id is not assigned";
    thread_context_ = std::make_unique<WorkerThreadContext>(current_job_id_);
  }

  return *thread_context_;
}
}  // namespace core
}  // namespace ray
